import math

import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse

from PIL import Image
import torch


# 加载图像
image_path = r'/root/autodl-tmp/Culane/driver_100_30frame/05251517_0433.MP4/00000.jpg'
image = Image.open(image_path)

# # 您的数据
data1 = [
    3.75226, 520.00000, 10.84336, 518.00000, 18.00389, 516.00000, 25.20584, 514.00000, 32.42116, 512.00000, 39.62185,
    510.00000,
    46.78102, 508.00000, 53.90847, 506.00000, 61.05715, 504.00000, 68.25561, 502.00000, 75.46951, 500.00000, 82.66186,
    498.00000,
    89.82703, 496.00000, 96.97119, 494.00000, 104.12436, 492.00000, 111.33024, 490.00000, 118.58790, 488.00000,
    125.84155,
    486.00000, 133.05399, 484.00000, 140.24149, 482.00000, 147.42549, 480.00000, 154.61277, 478.00000, 161.80550,
    476.00000,
    169.00686, 474.00000, 176.22024, 472.00000, 183.43233, 470.00000, 190.60873, 468.00000, 197.73150, 466.00000,
    204.82794,
    464.00000, 211.93029, 462.00000, 219.05950, 460.00000, 226.23260, 458.00000, 233.46501, 456.00000, 240.76989,
    454.00000,
    248.11336, 452.00000, 255.39961, 450.00000, 262.57584, 448.00000, 269.71402, 446.00000, 276.89270, 444.00000,
    284.11309,
    442.00000, 291.35043, 440.00000, 298.59310, 438.00000, 305.83908, 436.00000, 313.09856, 434.00000, 320.39813,
    432.00000,
    327.73994, 430.00000, 335.05572, 428.00000, 342.28145, 426.00000, 349.44888, 424.00000, 356.62110, 422.00000,
    363.81307,
    420.00000, 371.00480, 418.00000, 378.19036, 416.00000, 385.38490, 414.00000, 392.59879, 412.00000, 399.82413,
    410.00000,
    407.05002, 408.00000, 414.26696, 406.00000, 421.46681, 404.00000, 428.65386, 402.00000, 435.84147, 400.00000,
    443.02800,
    398.00000, 450.18926, 396.00000, 457.31380, 394.00000, 464.43351, 392.00000, 471.58385, 390.00000, 478.77158,
    388.00000,
    485.99149, 386.00000, 493.22834, 384.00000, 500.45921, 382.00000, 507.67222, 380.00000, 514.87268, 378.00000,
    522.06743,
    376.00000, 529.26357, 374.00000, 536.46571, 372.00000, 543.66164, 370.00000, 550.83477, 368.00000, 558.01005,
    366.00000,
    565.24554, 364.00000, 572.54456, 362.00000, 579.82186, 360.00000, 587.02174, 358.00000, 594.20983, 356.00000,
    601.46069,
    354.00000, 608.73962, 352.00000, 615.96967, 350.00000, 623.15199, 348.00000, 630.35340, 346.00000, 637.59831,
    344.00000,
    644.83764, 342.00000, 652.02873, 340.00000, 659.17383, 338.00000, 666.28658, 336.00000, 673.39466, 334.00000,
    680.53011,
    332.00000, 687.69154, 330.00000, 694.84844, 328.00000, 701.99300, 326.00000, 709.15811, 324.00000, 716.37228,
    322.00000,
    723.63459, 320.00000, 730.93310, 318.00000, 738.22140, 316.00000, 745.44017, 314.00000, 752.59048, 312.00000,
    759.72855,
    310.00000, 766.91263, 308.00000, 774.20094, 306.00000, 781.65173, 304.00000, 789.32324, 302.00000
]
data2 = [
    0, 520,
    41, 521,
    82, 522,
    123, 523,
    164, 524,
    205, 525,
    246, 526,
    287, 527,
    328, 528,
    369, 529,
    410, 530,
    451, 531,
    492, 532,
    533, 533,
    574, 534,
    615, 535,
    656, 536,
    697, 537,
    738, 538,
    779, 539,
    820, 540,
    861, 515,
    902, 490,
    943, 465,
    984, 440,
    1025, 415,
    1066, 390,
    1107, 365,
    1148, 340,
    1189, 315,
    1230, 290,
    1271, 265,
    1312, 240,
    1353, 215,
    1394, 190,
    1435, 165,
    1476, 140,
    1517, 115,
    1558, 90,
    1599, 65,
    1640, 40,
]
data3 = [
    0, 520,
    41, 520,
    82, 520,
    123, 520,
    164, 520,
    205, 520,
    246, 520,
    287, 520,
    328, 520,
    369, 520,
    410, 520,
    451, 520,
    492, 520,
    533, 520,
    574, 520,
    615, 520,
    656, 520,
    697, 520,
    738, 520,
    779, 520,
    820, 520,
    861, 520,
    902, 520,
    943, 520,
    984, 520,
    1025, 520,
    1066, 520,
    1107, 520,
    1148, 520,
    1189, 520,
    1230, 520,
    1271, 520,
    1312, 520,
    1353, 520,
    1394, 520,
    1435, 520,
    1476, 520,
    1517, 520,
    1558, 520,
    1599, 520,
    1640, 520,
]

data = data1



# 将数据分为X和Y坐标
# 将数据分为X和Y坐标
def split_data(data):
    Xs = data[::2]
    Ys = data[1::2]
    return Xs, Ys


Xs, Ys = split_data(data)
strip_size = Ys[1] - Ys[0]

# 创建图形和轴
fig, ax = plt.subplots(figsize=(10, 5))

# 显示图像
# ax.imshow(image, extent=[0, image.width, -image.height, 0])

# 绘制原始车道线
def plot_lines(ax, Xs, Ys, color):
    for i in range(len(Xs) - 1):
        x_values = [Xs[i], Xs[i + 1]]
        y_values = [Ys[i] * -1, Ys[i + 1] * -1]
        ax.plot(x_values, y_values, color)

plot_lines(ax, Xs, Ys, 'b-')

# 计算高斯核******************************************************
thetas = [
    max(
        math.atan(i * strip_size / (Xs[i] - Xs[0] + 1e-5)) / math.pi, 
        1 - abs(math.atan(i * strip_size / (Xs[i] - Xs[0] + 1e-5)) / math.pi)
    )
    for i in range(1, len(Xs))
]
theta_far = sum(thetas) / len(thetas)

pred_high = strip_size * len(Ys) / math.sin(math.pi * theta_far)
# pred_wide = 7.5
pred_wide = 15


center_x,center_y = Xs[len(Xs)//2],-1 * Ys[len(Ys)//2]

theta = theta_far * math.pi   
theta = torch.tensor(theta)
w = torch.tensor(pred_high*0.5)
h = torch.tensor(pred_wide*0.5) 

# 计算μ
mu = torch.tensor([center_x,center_y])

# 计算R
cos_r = torch.cos(theta)
sin_r = torch.sin(theta)
R = torch.stack((cos_r, -sin_r, sin_r, cos_r), dim=-1).reshape(-1, 2, 2)

# 计算Λ
zero_tensor = torch.tensor(0.0)
Lambda = torch.stack((w, zero_tensor, zero_tensor, h), dim=-1).reshape(-1, 2, 2)
Lambda = Lambda.square()

# 计算Σ
sigma = R.bmm(Lambda).bmm(R.permute(0, 2, 1)).reshape((2, 2))

# 创建二维高斯分布
def plot_gaussian(ax, mean, cov_matrix):
    eigenvalues, eigenvectors = torch.linalg.eig(cov_matrix)
    lambda_ = torch.sqrt(eigenvalues.real)
    angle = torch.atan2(eigenvectors[0, 1].real, eigenvectors[0, 0].real)
    angle = torch.rad2deg(angle).item()
    ell = Ellipse(xy=mean,
                  width=lambda_[0].item()*2, height=lambda_[1].item()*2,
                  angle=angle,
                  edgecolor='r', fc='None', lw=2)
    ax.add_patch(ell)

# 绘制高斯分布
mean = [center_x, center_y]
plot_gaussian(ax, mean, sigma)

#******************************************************
def plot_rotated_rectangle(ax, center_x, center_y, width, height, angle):
    # 计算矩形的四个顶点
    rect = torch.tensor([
        [-width / 2, -height / 2],
        [width / 2, -height / 2],
        [width / 2, height / 2],
        [-width / 2, height / 2]
    ])
    
    # 旋转矩形
    rotation_matrix = torch.tensor([
        [torch.cos(angle), -torch.sin(angle)],
        [torch.sin(angle), torch.cos(angle)]
    ])
    rotated_rect = torch.matmul(rect, rotation_matrix.T)
    
    # 平移矩形到中心点
    rotated_rect += torch.tensor([center_x, center_y])
    
    # 创建多边形
    polygon = plt.Polygon(rotated_rect.numpy(), edgecolor='r', facecolor='none')
    ax.add_patch(polygon)

plot_rotated_rectangle(ax, center_x, center_y, w * 2, h, -1 * theta)




ax.set_title('Image with Line Segments and 2D Gaussian Kernel')
ax.set_xlabel('X axis')
ax.set_ylabel('Y axis')

# 显示图表
plt.show()
plt.show()